{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "44310256-73be-4e34-8996-54eb3bfae52e",
   "metadata": {},
   "source": [
    "# Text Rerank with Qwen3 and OpenVINO\n",
    "\n",
    "The Qwen3 Embedding model series is the latest proprietary model of the Qwen family, specifically designed for text embedding and ranking tasks. Building upon the dense foundational models of the Qwen3 series, it provides a comprehensive range of text embeddings and reranking models in various sizes (0.6B, 4B, and 8B). This series inherits the exceptional multilingual capabilities, long-text understanding, and reasoning skills of its foundational model. The Qwen3 Embedding series represents significant advancements in multiple text embedding and ranking tasks, including text retrieval, code retrieval, text classification, text clustering, and bitext mining.\n",
    "\n",
    "<img src=\"https://qianwen-res.oss-accelerate.aliyuncs.com/logo_qwen_embedding.png\" width=\"400\"/>\n",
    "\n",
    "In this tutorial we consider how to convert and optimize Qwen3 Reranker model using OpenVINO.\n",
    "\n",
    "#### Table of contents:\n",
    "\n",
    "- [Prerequisites](#Prerequisites)\n",
    "- [Select model](#Select-model)\n",
    "- [Convert model using Optimum Intel](#Convert-model-using-Optimum-Intel)\n",
    "- [Run OpenVINO model inference with Optimum-intel](#Run-OpenVINO-model-inference-with-Optimum-intel)\n",
    "\n",
    "\n",
    "<img referrerpolicy=\"no-referrer-when-downgrade\" src=\"https://static.scarf.sh/a.png?x-pxid=5b5a4db0-7875-4bfb-bdbd-01698b5b1a77&file=notebooks/flux.1-image-generation/flux.1-image-generation.ipynb\" />\n",
    "\n",
    "### Installation Instructions\n",
    "\n",
    "This is a self-contained example that relies solely on its own code.\n",
    "\n",
    "We recommend  running the notebook in a virtual environment. You only need a Jupyter server to start.\n",
    "For details, please refer to [Installation Guide](https://github.com/openvinotoolkit/openvino_notebooks/blob/latest/README.md#-installation-guide).\n",
    "\n",
    "<img referrerpolicy=\"no-referrer-when-downgrade\" src=\"https://static.scarf.sh/a.png?x-pxid=5b5a4db0-7875-4bfb-bdbd-01698b5b1a77&file=notebooks/qwen3-embedding/qwen3-reranker.ipynb\" />\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ca26860-bbb7-47e9-9c56-a5e071853b81",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "[back to top ⬆️](#Table-of-contents:)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8699810f-7970-4b00-8ee3-30b8f72ab116",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m25.1.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m25.2\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
      "Note: you may need to restart the kernel to use updated packages.\n",
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m25.1.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m25.2\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "import platform\n",
    "\n",
    "%pip uninstall -q -y optimum optimum-intel optimum-onnx\n",
    "%pip install -q \"git+https://github.com/huggingface/optimum-intel.git\"\n",
    "%pip install -qU \"openvino>=2025.2\" \"openvino_tokenizers>=2025.2\"\n",
    "\n",
    "if platform.system() == \"Darwin\":\n",
    "    %pip install -q \"numpy<2.0.0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e44ef0c1-5ee8-4600-bb0c-847342cb42cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "from pathlib import Path\n",
    "\n",
    "if not Path(\"cmd_helper.py\").exists():\n",
    "    r = requests.get(url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/cmd_helper.py\")\n",
    "    open(\"cmd_helper.py\", \"w\").write(r.text)\n",
    "\n",
    "if not Path(\"notebook_utils.py\").exists():\n",
    "    r = requests.get(url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py\")\n",
    "    open(\"notebook_utils.py\", \"w\").write(r.text)\n",
    "\n",
    "# Read more about telemetry collection at https://github.com/openvinotoolkit/openvino_notebooks?tab=readme-ov-file#-telemetry\n",
    "from notebook_utils import collect_telemetry\n",
    "\n",
    "collect_telemetry(\"flux.1-image-generation.ipynb\")\n",
    "\n",
    "# Read more about telemetry collection at https://github.com/openvinotoolkit/openvino_notebooks?tab=readme-ov-file#-telemetry\n",
    "from notebook_utils import collect_telemetry\n",
    "\n",
    "collect_telemetry(\"qwen3-reranker.ipynb\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcecf123-3404-45be-8144-f1ff2bc55de2",
   "metadata": {},
   "source": [
    "## Select model\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Qwen3 Reranker Series Model list:\n",
    "\n",
    "| Model Type       | Models               | Size | Layers | Sequence Length | Embedding Dimension | MRL Support | Instruction Aware |\n",
    "|------------------|----------------------|------|--------|-----------------|---------------------|-------------|----------------|\n",
    "| Text Reranking   | [Qwen3-Reranker-0.6B](https://huggingface.co/Qwen/Qwen3-Reranker-0.6B) | 0.6B | 28     | 32K             | -                   | -           | Yes            |\n",
    "| Text Reranking   | [Qwen3-Reranker-4B](https://huggingface.co/Qwen/Qwen3-Reranker-4B)   | 4B   | 36     | 32K             | -                   | -           | Yes            |\n",
    "| Text Reranking   | [Qwen3-Reranker-8B](https://huggingface.co/Qwen/Qwen3-Reranker-8B)   | 8B   | 36     | 32K             | -                   | -           | Yes            |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d41261dc-ed3a-42d6-ab48-1736b6a8c245",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "197e13d1f99347269536116613c077b0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Reranker Model:', options=('Qwen/Qwen3-Reranker-0.6B', 'Qwen/Qwen3-Reranker-4B', 'Qwen/Q…"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import ipywidgets as widgets\n",
    "\n",
    "model_ids = [\"Qwen/Qwen3-Reranker-0.6B\", \"Qwen/Qwen3-Reranker-4B\", \"Qwen/Qwen3-Reranker-8B\"]\n",
    "\n",
    "model_selector = widgets.Dropdown(\n",
    "    options=model_ids,\n",
    "    default=model_ids[0],\n",
    "    description=\"Reranker Model:\",\n",
    ")\n",
    "\n",
    "model_selector"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7075f6e5-864c-49e4-ac92-cf0d6be56bf3",
   "metadata": {},
   "source": [
    "## Convert model using Optimum Intel\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "For convenience, we will use OpenVINO integration with HuggingFace Optimum. 🤗 [Optimum Intel](https://huggingface.co/docs/optimum/intel/index) is the interface between the 🤗 Transformers and Diffusers libraries and the different tools and libraries provided by Intel to accelerate end-to-end pipelines on Intel architectures.\n",
    "\n",
    "Among other use cases, Optimum Intel provides a simple interface to optimize your Transformers and Diffusers models, convert them to the OpenVINO Intermediate Representation (IR) format and run inference using OpenVINO Runtime. `optimum-cli` provides command line interface for model conversion and optimization. \n",
    "\n",
    "General command format:\n",
    "\n",
    "```bash\n",
    "optimum-cli export openvino --model <model_id_or_path> --task <task> <output_dir>\n",
    "```\n",
    "\n",
    "where task is task to export the model for, if not specified, the task will be auto-inferred based on the model. You can find a mapping between tasks and model classes in Optimum TaskManager [documentation](https://huggingface.co/docs/optimum/exporters/task_manager). Additionally, you can specify weights compression using `--weight-format` argument with one of following options: `fp32`, `fp16`, `int8` and `int4`. Fro int8 and int4 [nncf](https://github.com/openvinotoolkit/nncf) will be used for  weight compression. More details about model export provided in [Optimum Intel documentation](https://huggingface.co/docs/optimum/intel/openvino/export#export-your-model)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c3e3fb2d-5aa1-4765-b518-c147d175910a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "990aba75a12b46648c8d0ede5b9078a2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Checkbox(value=False, description='Weight compression'),))"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "to_compress = widgets.Checkbox(\n",
    "    value=False,\n",
    "    description=\"Weight compression\",\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "visible_widgets = [to_compress]\n",
    "\n",
    "options = widgets.VBox(visible_widgets)\n",
    "\n",
    "options"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94c966f8-30c6-4848-829e-9dfe88d5dd94",
   "metadata": {},
   "source": [
    "The Qwen3-embedding model can be exported by `text-generation` task with Optimum-intel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4c93db93-0a67-4dae-9caa-c2a4f399dce0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "model_id = model_selector.value\n",
    "\n",
    "model_base_dir = Path(model_id.split(\"/\")[-1])\n",
    "additional_args = {\"task\": \"text-generation\"}\n",
    "\n",
    "if to_compress.value:\n",
    "    model_dir = model_base_dir / \"INT8\"\n",
    "    additional_args.update({\"weight-format\": \"int8\"})\n",
    "else:\n",
    "    model_dir = model_base_dir / \"FP16\"\n",
    "    additional_args.update({\"weight-format\": \"fp16\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "63d78398-e94d-48fa-a834-3719a7bc7966",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "**Export command:**"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "`optimum-cli export openvino --model Qwen/Qwen3-Reranker-0.6B Qwen3-Reranker-0.6B/FP16 --task text-generation --weight-format fp16`"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from cmd_helper import optimum_cli\n",
    "\n",
    "if not model_dir.exists():\n",
    "    optimum_cli(model_id, model_dir, additional_args=additional_args)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e614f2c9-7ddc-41c7-bd19-b404d37137b8",
   "metadata": {},
   "source": [
    "## Run OpenVINO model inference with Optimum-intel\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Select device from dropdown list for running inference using OpenVINO."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "24e0ac3c-417d-480f-9a36-fbb48bb3f44e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f0c7990b9e0444698aac4e6c86b89805",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Device:', options=('CPU', 'GPU', 'AUTO'), value='CPU')"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from notebook_utils import device_widget\n",
    "\n",
    "device = device_widget(default=\"CPU\", exclude=[\"NPU\"])\n",
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b2af1c4-ff86-4fe2-8c04-3bf6b24bc284",
   "metadata": {},
   "source": [
    "The Qwen3-embedding model can be exported by class `OVModelForCausalLM` with Optimum-intel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "dc6b9f22-cc20-44c7-9fc8-b691b0f2c70b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from optimum.intel import OVModelForCausalLM\n",
    "\n",
    "model = OVModelForCausalLM.from_pretrained(model_dir, device=device.value, use_cache=False, export=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "98bcc2a5-8891-4159-80cd-0f5e27e4ebce",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You're using a Qwen2TokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n",
      "/home/ethan/intel/openvino_notebooks/openvino_env/lib/python3.10/site-packages/transformers/tokenization_utils_base.py:2696: UserWarning: `max_length` is ignored when `padding`=`True` and there is no truncation strategy. To pad to max length, use `padding='max_length'`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "scores:  [0.9994981288909912, 0.9993618130683899]\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import AutoTokenizer\n",
    "\n",
    "\n",
    "def format_instruction(instruction, query, doc):\n",
    "    if instruction is None:\n",
    "        instruction = \"Given a web search query, retrieve relevant passages that answer the query\"\n",
    "    output = \"<Instruct>: {instruction}\\n<Query>: {query}\\n<Document>: {doc}\".format(instruction=instruction, query=query, doc=doc)\n",
    "    return output\n",
    "\n",
    "\n",
    "def process_inputs(pairs):\n",
    "    inputs = tokenizer(\n",
    "        pairs, padding=False, truncation=\"longest_first\", return_attention_mask=False, max_length=max_length - len(prefix_tokens) - len(suffix_tokens)\n",
    "    )\n",
    "    for i, ele in enumerate(inputs[\"input_ids\"]):\n",
    "        inputs[\"input_ids\"][i] = prefix_tokens + ele + suffix_tokens\n",
    "    inputs = tokenizer.pad(inputs, padding=True, return_tensors=\"pt\", max_length=max_length)\n",
    "    for key in inputs:\n",
    "        inputs[key] = inputs[key].to(model.device)\n",
    "    return inputs\n",
    "\n",
    "\n",
    "def compute_logits(inputs, **kwargs):\n",
    "    batch_scores = model(**inputs).logits[:, -1, :]\n",
    "    true_vector = batch_scores[:, token_true_id]\n",
    "    false_vector = batch_scores[:, token_false_id]\n",
    "    batch_scores = torch.stack([false_vector, true_vector], dim=1)\n",
    "    batch_scores = torch.nn.functional.log_softmax(batch_scores, dim=1)\n",
    "    scores = batch_scores[:, 1].exp().tolist()\n",
    "    return scores\n",
    "\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_dir, padding_side=\"left\")\n",
    "token_false_id = tokenizer.convert_tokens_to_ids(\"no\")\n",
    "token_true_id = tokenizer.convert_tokens_to_ids(\"yes\")\n",
    "max_length = 8192\n",
    "\n",
    "prefix = '<|im_start|>system\\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be \"yes\" or \"no\".<|im_end|>\\n<|im_start|>user\\n'\n",
    "suffix = \"<|im_end|>\\n<|im_start|>assistant\\n<think>\\n\\n</think>\\n\\n\"\n",
    "prefix_tokens = tokenizer.encode(prefix, add_special_tokens=False)\n",
    "suffix_tokens = tokenizer.encode(suffix, add_special_tokens=False)\n",
    "\n",
    "task = \"Given a web search query, retrieve relevant passages that answer the query\"\n",
    "\n",
    "queries = [\n",
    "    \"What is the capital of China?\",\n",
    "    \"Explain gravity\",\n",
    "]\n",
    "\n",
    "documents = [\n",
    "    \"The capital of China is Beijing.\",\n",
    "    \"Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun.\",\n",
    "]\n",
    "\n",
    "pairs = [format_instruction(task, query, doc) for query, doc in zip(queries, documents)]\n",
    "\n",
    "# Tokenize the input texts\n",
    "inputs = process_inputs(pairs)\n",
    "scores = compute_logits(inputs)\n",
    "\n",
    "print(\"scores: \", scores)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  },
  "openvino_notebooks": {
   "imageUrl": "https://qianwen-res.oss-accelerate.aliyuncs.com/logo_qwen_embedding.png",
   "tags": {
    "categories": [
     "Model Demos",
     "AI Trends"
    ],
    "libraries": [],
    "other": [],
    "tasks": [
     "Feature Extraction"
    ]
   }
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
